/*
 * Javlov - a Java toolkit for reinforcement learning with multi-agent support.
 * 
 * Copyright (c) 2009 Matthijs Snel
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.policy;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

import net.javlov.Option;
import net.javlov.Policy;
import net.javlov.QFunction;
import net.javlov.State;
import net.javlov.util.ArrayUtil;

/**
 * Picks actions from Q-values according to a softmax distribution over the values:
 * 
 * {@code p(s,a) = exp(Q(s,a)) / sum_all_actions( exp(Q(s,a)) )}.
 * 
 * @author Matthijs Snel
 *
 */
//TODO code duplication in getOption and getOptionProbabilities
public class SoftmaxPolicy implements Policy {

	/**
	 * The Q-value function.
	 */
	private QFunction q;
	
	/**
	 * List of allowed actions. Index of actions in the list should correspond to their
	 * ID.
	 */
	protected List<? extends Option> optionPool;
	
	/**
	 * Random number generator.
	 */
	private Random rng;
	
	public SoftmaxPolicy(QFunction qf, List<? extends Option> options) {
		setQFunction(qf);
		optionPool = options;
		rng = new Random();
	}
	
	public void setQFunction(QFunction q) {
		this.q = q;
	}

	public QFunction getQFunction() {
		return q;
	}
	
	public <T> double[] getOptionProbabilities( State<T> s, double[] qvalues ) {
		List<? extends Option> stateOptionSet = getStateOptionSet(s);
		if ( stateOptionSet == null || stateOptionSet.size() == qvalues.length )
			return getSoftmaxValues(qvalues);

		double[] qvals = new double[stateOptionSet.size()];
		for ( int i = 0; i < qvals.length; i++ )
			qvals[i] = qvalues[stateOptionSet.get(i).getID()];
		qvals = getSoftmaxValues(qvals);
		
		int i = 0;
		double[] probs = new double[qvalues.length];
		for ( Option opt : stateOptionSet )
			probs[opt.getID()] = qvals[i++];
		return probs;
	}
	
	public static double[] getSoftmaxValues(double[] values) {
		double	r[] = new double[values.length],
				exp[] = new double[values.length],
				sum = 0;
		for ( int i = 0; i < values.length; i++ ) {
			exp[i] = Math.exp(values[i]);
			sum += exp[i];
		}
		for ( int i = 0; i < values.length; i++ )
			r[i] = exp[i] / sum;
		return r;
	}
	
	protected Option pickSoftmaxOption(List<? extends Option> options, double[] qvalues) {
		double softmaxVals[] = getSoftmaxValues(qvalues);
		//create a "ladder" of the values where each entry is added to the previous one
		ArrayUtil.sumeachInPlace(softmaxVals);
		double r = rng.nextDouble();
		int i;
		for ( i = 0; i < softmaxVals.length; i++ )
			if (r < softmaxVals[i])
				break;
		return options.get(i);
	}

	@Override
	public <T> Option getOption(State<T> s) {
		return getOption( s, q.getValues(s) );
	}

	protected <T> List<Option> getStateOptionSet(State<T> s) {
		List<Option> stateOptionSet = new ArrayList<Option>(optionPool.size());
		for ( Option o : optionPool )
			if ( o.isEligible(s) )
				stateOptionSet.add(o);
		
		if ( stateOptionSet.size() == 0 )
			throw new RuntimeException("No eligible options for state: " + s);
		
		return stateOptionSet;
	}
	
	@Override
	public <T> Option getOption(State<T> s, double[] qvalues) {
		List<? extends Option> stateOptionSet = getStateOptionSet(s);
		if ( stateOptionSet == null || stateOptionSet.size() == qvalues.length ) {
			return pickSoftmaxOption(optionPool, qvalues);
		}
		else {
			double[] qvals = new double[stateOptionSet.size()];
			for ( int i = 0; i < qvals.length; i++ )
				qvals[i] = qvalues[stateOptionSet.get(i).getID()];
			return pickSoftmaxOption(stateOptionSet, qvals);
		}
	}

	@Override
	public void init() {
		for ( Option o : optionPool )
			o.init();
	}

	@Override
	public void reset() {
		for ( Option o : optionPool )
			o.reset();
	}

}
